import math as m
import numpy as np
norm = np.linalg.norm

def main():
    fin=open("0.lammpstrj")
    fout=open("a.xyz", "w+")
    natom = 22
    xyz = np.zeros([natom, 3])
    vxyz = np.zeros([natom, 3])
    readxyz(fin, xyz)

def read_energy(name):
    e = np.loadtxt("{}.energy".format(name))
    return e

def read_intc(name):
    col = np.loadtxt("{}.intc".format(name))
    if (len(col.shape)>1):
        return col[:, 1:]
    else:
        return col[1:].reshape([1, -1])

def read_trj(name):
    try:
        with open("{}.lammpstrj".format(name)) as fin:
            fin.readline()
            fin.readline()
            fin.readline()
            a = fin.readline()
            natom = int(a.split()[0])
    except:
        raise NameError("the file cannot be read")

    xyz = np.zeros([natom, 3])
    allxyz = []
    with open("{}.lammpstrj".format(name)) as fin:
        while (readxyz(fin, xyz)):
            allxyz += [np.array(xyz.flat)]

    current_xyz = np.array(allxyz)

    return current_xyz

def readxyz(fin, xyz):
    a = fin.readline()
    if (a):
      fin.readline()
      fin.readline()
      a = fin.readline()
      natom = int(a.split()[0])
      fin.readline()
      fin.readline()
      fin.readline()
      fin.readline()
      fin.readline()
      for i in range(natom):
          line = fin.readline().split()
          index = int(line[0])-1
          xyz[index, :] = np.array(list(map(float, line[1:4])))
      # xyz = rotateCtop(natom, xyz, 8, 10, 6)
      return True
    else:
        return False


def rotateCtop(natom, xyz0, centerid, topid, leftid):

    xyz = np.copy(xyz0)
    xyz = xyz - xyz[centerid, :]
    txyz = xyz[topid,:]
    dr = txyz
    r_tc = norm(dr)
    r_tcxy = norm(dr[:2])
    ct = dr[2]/r_tc
    st = m.sqrt(1-ct**2)
    cp = dr[0]/r_tcxy
    sp = dr[1]/r_tcxy
    rotation_matrix = np.array([
        [ct*cp, ct*sp, -st],
        [-sp, cp, 0],
        [st*cp, st*sp, ct]]).T
    # all xyz is operated by it
    xyz = xyz.dot(rotation_matrix)

    leftxyz = xyz[leftid,:]

    dr = leftxyz
    r_lcxy = norm(dr[:2])
    cp = dr[0]/r_lcxy
    sp = dr[1]/r_lcxy
    rotation_matrix = np.array([
        [cp, sp, 0],
        [-sp, cp, 0],
        [0, 0, 1.0]]).T
    # all xyz is operated by it
    xyz = xyz.dot(rotation_matrix)
    return xyz

if __name__ == '__main__':
    main()
